#!/usr/bin/env python3
"""
将 AlgViz 采集到的 NDJSON 事件整理为结构化 JSON，并输出统计信息。

用法示例：
    python scripts/postprocess_ndjson.py snapshots.ndjson \
        --output outputs/backtrack_trace.json --pretty
"""

from __future__ import annotations

import argparse
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, List


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="AlgViz NDJSON 后处理脚本",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "input",
        type=Path,
        help="NDJSON 事件文件路径（例如 snapshots.ndjson）",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("outputs/trace.json"),
        help="输出 JSON 文件路径",
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="以缩进格式写入输出 JSON，便于人工检查",
    )
    return parser.parse_args()


def load_events(path: Path) -> List[Dict[str, Any]]:
    events: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as fh:
        for line_no, line in enumerate(fh, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                events.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise ValueError(f"第 {line_no} 行解析失败: {exc}") from exc
    return events


def extract_solutions(events: Iterable[Dict[str, Any]]) -> List[Dict[str, Any]]:
    solutions: List[Dict[str, Any]] = []
    for event in events:
        if event.get("event") != "solution":
            continue
        payload = {
            "timestamp": event.get("timestamp"),
            "frame": event.get("frame"),
            "location": event.get("location"),
            "path": event.get("variables", {}).get("path"),
            "target": event.get("variables", {}).get("target"),
            "depth": event.get("depth"),
        }
        solutions.append(payload)
    return solutions


def build_stats(events: Iterable[Dict[str, Any]]) -> Dict[str, Any]:
    event_counter = Counter()
    frame_counter = Counter()
    variable_presence = defaultdict(int)

    for event in events:
        event_counter[event.get("event")] += 1
        frame_counter[event.get("frame")] += 1
        for var_name, value in event.get("variables", {}).items():
            if value is not None:
                variable_presence[var_name] += 1

    stats = {
        "total_events": sum(event_counter.values()),
        "events_by_type": event_counter,
        "events_by_frame": frame_counter,
        "variable_presence": variable_presence,
    }

    # 将 Counter/defaultdict 转成普通 dict，便于 JSON 序列化
    return {
        key: dict(value) if isinstance(value, (Counter, defaultdict)) else value
        for key, value in stats.items()
    }


def normalize_events(events: Iterable[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """确保输出中的事件字段顺序统一、变量为空时以 None 表示。"""
    normalized: List[Dict[str, Any]] = []
    for event in events:
        variables = event.get("variables", {})
        clean_variables = {
            name: (value if value is not None else None)
            for name, value in variables.items()
        }
        normalized.append(
            {
                "timestamp": event.get("timestamp"),
                "event": event.get("event"),
                "frame": event.get("frame"),
                "location": event.get("location"),
                "depth": event.get("depth"),
                "variables": clean_variables,
            }
        )
    return normalized


def main() -> None:
    args = parse_args()

    if not args.input.exists():
        raise FileNotFoundError(f"输入文件不存在：{args.input}")

    events = load_events(args.input)
    solutions = extract_solutions(events)
    stats = build_stats(events)
    normalized_events = normalize_events(events)

    result = {
        "metadata": {
            "source": str(args.input),
            "total_solutions": len(solutions),
        },
        "stats": stats,
        "solutions": solutions,
        "events": normalized_events,
    }

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as fh:
        json.dump(result, fh, ensure_ascii=False, indent=2 if args.pretty else None)

    print(f"[postprocess] 共读取 {len(events)} 条事件，提取 {len(solutions)} 个解。")
    print(f"[postprocess] 统计信息写入：{args.output}")


if __name__ == "__main__":
    main()
